package liming.ai.tongyi;

import com.openai.client.OpenAIClient;
import com.openai.core.JsonNull;
import com.openai.core.JsonValue;
import com.openai.models.*;
import liming.ai.abstracts.AbstractNet;
import liming.ai.abstracts.MessageContext;
import liming.tool.json.FileConfig;
import liming.tool.json.ValueOf;
import liming.tool.json.annotation.JsonMark;
import liming.tool.net.NetServer;
import liming.tool.runnable.HandRunnable;
import liming.tool.runnable.HaveReturnValueInterface;
import org.json.JSONObject;

import java.util.*;
import java.util.stream.Stream;

public class TongYiNet extends AbstractNet {
    OpenAIClient client;
    TongYiConfig config;
    public TongYiNet() {
        super();
        client = NetServer.OpenAIConfig.builder().build();
        config = new TongYiConfig();
        config.deserialize(FileConfig.getConfig(config.getJsonRootName()));
    }
    @Override
    public ResponseMessage chat(String query, MessageContext messageContext) {
        messageContext.addMessage(new TongYiMessage(TongYiMessage.Role.USER,query));
        List<ChatCompletionMessageParam> messages=new ArrayList<>();
        for (JSONObject message : messageContext.getMessages()) {
            String role = message.optString("role");
            String content = message.optString("content");

            switch (role) {
                case "system":
                    messages.add(ChatCompletionMessageParam.ofSystem(ChatCompletionSystemMessageParam.builder()
                            .content(ValueOf.valueOf(String.class,content))
                            .build()));
                    break;
                case "assistant":
                    messages.add(ChatCompletionMessageParam.ofAssistant(ChatCompletionAssistantMessageParam.builder()
                            .content(ValueOf.valueOf(String.class,content))
                            .build()));
                    break;
                case "user":
                    messages.add(ChatCompletionMessageParam.ofUser(ChatCompletionUserMessageParam.builder()
                            .content(ValueOf.valueOf(String.class,content))
                            .build()));
            }
        }
        ChatCompletionCreateParams.Builder params = ChatCompletionCreateParams.builder()
                .messages(messages)
                .model(config.body.model)
                .temperature(config.body.temperature)
                .topP(config.body.top_p)
                .presencePenalty(config.body.presence_penalty);
        Map<String ,JsonValue> data = new HashMap<>();
        data.put("stream", JsonValue.from(config.body.stream));
        data.put("enable_search", JsonValue.from(config.body.enable_search));
        data.put("enable_thinking", JsonValue.from(config.body.enable_thinking));
        params.additionalBodyProperties(data);

        TongYiResponseMessage message;
        // 流式处理分支
        if (Boolean.TRUE.equals(config.body.stream)) {
            message = handleStreamResponse(client.chat().completions().createStreaming(params.build()).stream());
        } else {
            message =handleNormalResponse(client.chat().completions().create(params.build()));
        }
        if(message.status==200)
            messageContext.addMessage(new TongYiMessage(TongYiMessage.Role.ASSISTANT,message.body));
        else
            messageContext.removeLastMessage();
        return message;
    }

    // 新增私有方法处理响应
    private TongYiResponseMessage handleNormalResponse(ChatCompletion chatCompletion) {
        return HandRunnable.run(new HaveReturnValueInterface<TongYiResponseMessage>() {
            @Override
            public TongYiResponseMessage run() throws Exception {
                String response = chatCompletion.choices().stream()
                        .findFirst()
                        .map(choice -> choice.message().content().orElse("无返回内容"))
                        .orElse("无响应内容");
                return new TongYiResponseMessage(200, "Success", response);
            }
            @Override
            public TongYiResponseMessage onException(Exception e) {
                return new TongYiResponseMessage(500, e.getMessage(), null,e);
            }
        });
    }

    private TongYiResponseMessage handleStreamResponse(Stream<ChatCompletionChunk> stream) {
        return HandRunnable.run(new HaveReturnValueInterface<TongYiResponseMessage>() {
            @Override
            public TongYiResponseMessage run() throws Exception {
                TongYiResponseMessage success = new TongYiResponseMessage(200, "Success", null);
                stream.forEach(chunk -> chunk.choices().forEach(choice -> {
                    success.addBody(choice.delta().content().orElse(""));
                    success.addThinking(choice.delta()._additionalProperties().get("reasoning_content"));
                }));
                return success;
            }

            @Override
            public TongYiResponseMessage onException(Exception e) {
                return new TongYiResponseMessage(500, e.getMessage(), "","", e);
            }
        });
    }

    public static class TongYiResponseMessage extends ResponseMessage {
        @JsonMark
        public int status;
        @JsonMark
        public String message;
        @JsonMark
        public String thinking;
        @JsonMark
        public String  body;
        @JsonMark
        public Throwable throwable;
        public TongYiResponseMessage(){}
        public TongYiResponseMessage(int status, String message, String body){
            this(status, message, body, "");
        }
        public TongYiResponseMessage(int status, String message, String body,String thinking){
            this(status, message, body,thinking, null);
        }
        public TongYiResponseMessage(int status, String message, String body,Throwable throwable){
            this(status, message, body,null, throwable);
        }
        public TongYiResponseMessage(int status, String message, String body,String thinking, Throwable throwable){
            this.status = status;
            this.message = message;
            this.body = body;
            this.thinking = thinking;
            this.throwable = throwable;
        }

        void addBody(String body){
            this.body=ValueOf.valueOf(String.class,this.body) +ValueOf.valueOf(String.class,body);
        }
        void addThinking(JsonValue thinking){
            if(thinking instanceof JsonNull|| thinking==null) return;
            this.thinking = ValueOf.valueOf(String.class,this.thinking) +ValueOf.valueOf(String.class,thinking.toString());
        }
    }
}
